(*
  Copyright (c) 2009 Barry Schwartz

  Permission is hereby granted, free of charge, to any person
  obtaining a copy of this software and associated documentation
  files (the "Software"), to deal in the Software without
  restriction, including without limitation the rights to use,
  copy, modify, merge, publish, distribute, sublicense, and/or sell
  copies of the Software, and to permit persons to whom the
  Software is furnished to do so, subject to the following
  conditions:

  The above copyright notice and this permission notice shall be
  included in all copies or substantial portions of the Software.

  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
  HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
  OTHER DEALINGS IN THE SOFTWARE.
*)

(* A dynamic trie |t 'a| implements a map |uc_string -> 'a| that supports adding new elements. *)

open UTypes
  
type 'a t = { children : ('a t) DynamicCharMap.t; data : 'a option }

let empty = { children = DynamicCharMap.empty; data = None; }
  
let is_empty trie =
  (trie.data = None) && (DynamicCharMap.is_empty trie.children)
  
let prefix trie char =
  try DynamicCharMap.find char trie.children with | Not_found -> empty
  
let root_value trie = trie.data
  
let generic_lookup fold str trie = root_value (fold prefix trie str)
  
(*
  For efficiency reasons the following definitions are expanded:

value lookup_string str trie = generic_lookup Array.fold_left str trie;
value lookup_list   str trie = generic_lookup List.fold_left  str trie;
*)
let lookup_string str trie =
  let rec iter i trie =
    if i >= (Array.length str)
    then trie.data
    else
      (try
         let next = DynamicCharMap.find str.(i) trie.children
         in iter (i + 1) next
       with | Not_found -> None)
  in iter 0 trie
  
let rec lookup_list str trie =
  match str with
  | [] -> trie.data
  | c :: cs ->
      (try
         let next = DynamicCharMap.find c trie.children
         in lookup_list cs next
       with | Not_found -> None)
  
(*
  |lookup_prefix <str> <trie>| returns the data associated to the longest prefix of <str> that
  is found in <trie>.
*)
let generic_lookup_prefix fold str trie =
  root_value
    (fst
       (fold
          (fun (last, t) c ->
             let next = prefix t c
             in
               match next.data with
               | Some _ -> (next, next)
               | None -> (last, next))
          (trie, trie) str))
  
(*
  For efficiency reasons the following definitions are expanded:

value lookup_prefix_string str trie = generic_lookup_prefix Array.fold_left str trie;
value lookup_prefix_list   str trie = generic_lookup_prefix List.fold_left  str trie;
*)
let lookup_prefix_string str trie =
  let rec iter i trie data =
    if i >= (Array.length str)
    then (match trie.data with | Some _ -> trie.data | None -> data)
    else
      (try
         let next = DynamicCharMap.find str.(i) trie.children
         in
           match next.data with
           | Some _ -> iter (i + 1) next next.data
           | None -> iter (i + 1) next data
       with | Not_found -> data)
  in iter 0 trie None
  
let lookup_prefix_list str trie =
  let rec iter trie str data =
    match str with
    | [] -> (match trie.data with | Some _ -> trie.data | None -> data)
    | c :: cs ->
        (try
           let next = DynamicCharMap.find c trie.children
           in
             match next.data with
             | Some _ -> iter next cs next.data
             | None -> iter next cs data
         with | Not_found -> data)
  in iter trie str None
  
(*
  |lookup_prefix_stream <is> <trie>| returns the data associated to the longest prefix of <is> that
  is found in <trie>. The matching prefix will be removed from <is>.
*)
let lookup_prefix_stream is trie =
  let stream = UCStream.duplicate is in
  let rec iter trie len pos data =
    if UCStream.eof stream
    then
      (match trie.data with
       | Some _ -> ((trie.data), pos)
       | None -> (data, len))
    else
      (try
         let next = DynamicCharMap.find (UCStream.pop stream) trie.children
         in
           match next.data with
           | Some _ -> iter next (pos + 1) (pos + 1) next.data
           | None -> iter next len (pos + 1) data
       with | Not_found -> (data, len)) in
  let (data, len) = iter trie 0 0 None in (UCStream.remove is len; data)
  
let find_string str trie =
  match lookup_string str trie with | Some x -> x | None -> raise Not_found
  
let mem_string str trie =
  match lookup_string str trie with | Some _ -> true | None -> false
  
let find_list str trie =
  match lookup_list str trie with | Some x -> x | None -> raise Not_found
  
let mem_list str trie =
  match lookup_list str trie with | Some _ -> true | None -> false
  
let add_string str data trie =
  let rec iter i trie =
    if i >= (Array.length str)
    then { (trie) with data = Some data; }
    else
      (let next =
         try DynamicCharMap.find str.(i) trie.children
         with | Not_found -> empty
       in
         {
           (trie)
           with
           children =
             DynamicCharMap.add str.(i) (iter (i + 1) next) trie.children;
         })
  in iter 0 trie
  
let remove_string str trie =
  let rec iter i trie =
    if i >= (Array.length str)
    then { (trie) with data = None; }
    else
      (try
         let next = DynamicCharMap.find str.(i) trie.children
         in
           {
             (trie)
             with
             children =
               DynamicCharMap.add str.(i) (iter (i + 1) next) trie.children;
           }
       with | Not_found -> trie)
  in iter 0 trie
  
let rec add_list str data trie =
  match str with
  | [] -> { (trie) with data = Some data; }
  | c :: cs ->
      let next =
        (try DynamicCharMap.find c trie.children with | Not_found -> empty)
      in
        {
          (trie)
          with
          children =
            DynamicCharMap.add c (add_list cs data next) trie.children;
        }
  
let rec remove_list str trie =
  match str with
  | [] -> { (trie) with data = None; }
  | c :: cs ->
      (try
         let next = DynamicCharMap.find c trie.children
         in
           {
             (trie)
             with
             children =
               DynamicCharMap.add c (remove_list cs next) trie.children;
           }
       with | Not_found -> trie)
  
let merge t1 t2 =
  {
    data = (match t2.data with | None -> t1.data | Some _ -> t2.data);
    children =
      if DynamicCharMap.is_empty t1.children
      then t2.children
      else
        DynamicCharMap.fold
          (fun c x children -> DynamicCharMap.add c x children) t2.children
          t1.children;
  }
  
let rec map f trie =
  {
    children = DynamicCharMap.map (map f) trie.children;
    data = (match trie.data with | Some x -> Some (f x) | None -> None);
  }
  
let mapi f trie =
  let rec iter str trie =
    {
      children =
        DynamicCharMap.mapi (fun c t -> iter (c :: str) t) trie.children;
      data =
        (match trie.data with
         | Some x -> Some (f (XList.rev_to_array str) x)
         | None -> None);
    }
  in iter [] trie
  
let rec iter f trie =
  let rec loop str trie =
    ((match trie.data with
      | Some x -> f (XList.rev_to_array str) x
      | None -> ());
     DynamicCharMap.iter (fun c t -> loop (c :: str) t) trie.children)
  in loop [] trie
  
let fold f trie e =
  let rec iter str trie e =
    let x =
      match trie.data with
      | Some y -> f (XList.rev_to_array str) y e
      | None -> e
    in DynamicCharMap.fold (fun c t y -> iter (c :: str) t y) trie.children x
  in iter [] trie e
  
let rec depth trie =
  if DynamicCharMap.is_empty trie.children
  then 0
  else
    (let d =
       DynamicCharMap.fold (fun _ t d -> max d (depth t)) trie.children 0
     in d + 1)
